package com.streamkar.common.es;

import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpEntity;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.entity.ContentType;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
import org.apache.http.nio.entity.NStringEntity;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * ES Rest客户端工具
 */
@Component
@EnableConfigurationProperties({EsProperties.class})
public class ESRestClient implements InitializingBean {

    private static final Logger logger = LoggerFactory.getLogger(ESRestClient.class);

    @Autowired
    EsProperties esProperties;

    RestHighLevelClient restClient = null;

    @Override
    public void afterPropertiesSet() throws Exception {

        String[] hostArray = esProperties.getHosts().split(",");

        if (hostArray.length == 0) {
            throw new Exception("Es rest client's hosts is wrong");
        }

        try {
            List<HttpHost> httpHosts = new ArrayList<HttpHost>();
            for (String host : hostArray) {
                String[] hostPort = host.split(":");
                HttpHost httpHost = new HttpHost(hostPort[0], Integer.parseInt(hostPort[1]), "http");
                httpHosts.add(httpHost);
            }
            HttpHost[] arrHttpHost = new HttpHost[hostArray.length];

            RestClientBuilder builder = RestClient.builder(httpHosts.toArray(arrHttpHost));
            final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
            credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(esProperties.getUserName(), esProperties.getPassword()));
            builder.setHttpClientConfigCallback(new RestClientBuilder.HttpClientConfigCallback() {
                @Override
                public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpClientBuilder) {
                    httpClientBuilder.disableAuthCaching();
                    return httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
                }
            });
            restClient = new RestHighLevelClient(builder);

        } catch (Exception e) {
            e.printStackTrace();
            logger.error("ES rest client init failed!", e);
        }
    }

    public RestHighLevelClient getRestClient() {
        return restClient;
    }

    public Response performRequest(String method, String endpoint, String dslJson) throws IOException {
        if (StringUtils.isEmpty(dslJson)) {
            logger.error("dsl json string is empty!");
            return null;
        }

        HttpEntity entity = new NStringEntity(dslJson, ContentType.APPLICATION_JSON);
        return restClient.getLowLevelClient().performRequest(method, endpoint, Collections.<String, String>emptyMap(), entity);
    }
}
